from typing import Tuple, Optional, List
from numpy.typing import DTypeLike
import cv2
import numpy as np
import numexpr as ne
from common.logger import logger
from common.geoimage.raster_dataset import RasterDataset

####################
# dictionary for storing instruction on how to process Sentinel-2 data for normalized index or single band extraction.
####################
S2_ndxi_parameters = {
    "ndvi": {"bands": ["B08", "B04"], "expr": "(a-b)/(a+b)", "params": ["a", "b"]},
    "bndvi": {"bands": ["B08", "B02"], "expr": "(a-b)/(a+b)", "params": ["a", "b"]},
    "gndvi": {"bands": ["B08", "B03"], "expr": "(a-b)/(a+b)", "params": ["a", "b"]},
    "gbndvi": {
        "bands": ["B08", "B03", "B02"],
        "expr": "(a-(b+c))/(a+(b+c))",
        "params": ["a", "b", "c"],
    },
    "grndvi": {
        "bands": ["B08", "B03", "B04"],
        "expr": "(a-(b+c))/(a+(b+c))",
        "params": ["a", "b", "c"],
    },
    "mndvi": {"bands": ["B08", "B11"], "expr": "(a-b)/(a+b)", "params": ["a", "b"]},
    "lswi": {"bands": ["B08", "B11"], "expr": "(a-b)/(a+b)", "params": ["a", "b"]},
    "wdrvi01": {
        "bands": ["B08", "B04"],
        "expr": "((0.1*a)-b)/((0.1*a)+b)",
        "params": ["a", "b"],
    },
    "wdrvi05": {
        "bands": ["B08", "B04"],
        "expr": "((0.5*a)-b)/((0.5*a)+b)",
        "params": ["a", "b"],
    },
    "wdrvi001": {
        "bands": ["B08", "B04"],
        "expr": "((0.01*a)-b)/((0.01*a)+b)",
        "params": ["a", "b"],
    },
    "gcvi": {"bands": ["B08", "B03"], "expr": "a/b - 1", "params": ["a", "b"]},
    "cvi": {
        "bands": ["B08", "B04", "B03"],
        "expr": "(a/c)*(b/c)",
        "params": ["a", "b", "c"],
    },
    "evi": {
        "bands": ["B08", "B04", "B02"],
        "expr": "2.5 * ((a - b) / ((a + (6.0*b) - (7.5*c)) + 1.0))",
        "params": ["a", "b", "c"],
    },
    "evi2": {
        "bands": ["B08", "B04"],
        "expr": "2.5 * ((a - b) / (a + (2.4*b) + 1.0))",
        "params": ["a", "b"],
    },
    "vari": {
        "bands": ["B03", "B04", "B02"],
        "expr": "(a-b)/(a+b-c)",
        "params": ["a", "b", "c"],
    },
    "gari": {
        "bands": ["B08", "B03", "B02", "B04"],
        "expr": "(a-(b-(c-d)))/(a-(b+(c-d)))",
        "params": ["a", "b", "c", "d"],
    },
    "ngrdi": {"bands": ["B03", "B04"], "expr": "(a-b)/(a+b)", "params": ["a", "b"]},
    "gli": {
        "bands": ["B03", "B04", "B02"],
        "expr": "((2*a)-b-c)/((2*a)+b+c)",
        "params": ["a", "b", "c"],
    },
    "gvmi": {
        "bands": ["B08", "B11"],
        "expr": "((a+0.1)-(b+0.02))/((a+0.1)+(b+0.02))",
        "params": ["a", "b"],
    },
    "nbr": {"bands": ["B8A", "B12"], "expr": "(a-b)/(a+b)", "params": ["a", "b"]},
    "nir": {"bands": ["B08"], "expr": "a", "params": ["a"]},
    "red": {"bands": ["B04"], "expr": "a", "params": ["a"]},
    "green": {"bands": ["B03"], "expr": "a", "params": ["a"]},
    "blue": {"bands": ["B02"], "expr": "a", "params": ["a"]},
    "rededge": {"bands": ["B05"], "expr": "a", "params": ["a"]},
    "nir1": {"bands": ["B06"], "expr": "a", "params": ["a"]},
    "nir2": {"bands": ["B07"], "expr": "a", "params": ["a"]},
    "swir1": {"bands": ["B11"], "expr": "a", "params": ["a"]},
    "swir2": {"bands": ["B12"], "expr": "a", "params": ["a"]},
    "ccci": {
        "bands": ["B8A", "B05", "B04"],
        "expr": "((a-b)/(a+b))/((a-c)/(a+c))",
        "params": ["a", "b", "c"],
    },
}

# collection of dictionary
ndxi_collection = {
    "Sentinel-2": S2_ndxi_parameters,
    "Landsat-8": None,
    "MODIS": None,
}
satellites_sr_valid_range_dict = {
    "Sentinel-2": [0, 10000, 10000],  # min, max, scaling factor
    "Landsat-8": None,
    "MODIS": None,
}
####################
# dictionary for storing instruction on processing Sentinel-2 mask.
####################
mask_dict = {
    "Sentinel-2": {
        "L2A": {
            "bands": ["SCL"],
            "mask_index": [0],
            "VALS_TO_MASK": [[0, 3, 8, 9, 11]],
            "RANGE_NOT_TO_MASK": [None],
        },
        "L3A": {
            "bands": ["SCL", "QA"],
            "mask_index": [0, 1],
            "VALS_TO_MASK": [
                [0, 3, 9, 10],
                None,
            ],  # see L3ASCLLabel in catalog/models/sentinel2_l3a.py
            "RANGE_NOT_TO_MASK": [
                None,
                [50, 100],
            ],  # see L3AQALevel in catalog/models/sentinel2_l3a.py # Also consider keeping only 100 in production.
        },
    },
    "Landsat-8": None,
    "MODIS": None,
}


def compute_ndxi(
    nda: np.ndarray,
    ndxi_target: str = "ndvi",
    working_dtype: DTypeLike = np.float32,
    satellite_name: str = "Sentinel-2",
    expr: Optional[str] = None,
    params: List[str] = None,
    nda_mask: Optional[np.ndarray] = None,
) -> Tuple[np.ndarray, np.ndarray]:
    """Compute metrics using expressions with "numexpr"
    NOTE: This method is completely rewritten, the original code is backed up at
    projects/20230111_yield_analytic_workflow_engineering_backup/derived_compute_ndxi_backup.py
    This numexpr computing method, according to Zhou Hang, is fast for large array,
    for example, working on 10980 X 10980 pixels; numexpr would be about twice as
    fast as the option of calculating using numpy only; However, it may be slower
    than numpy for small array.

    Parameters
    ----------
    nda : np.ndarray
        The input numpy array. Must be a 3D numpy array of with dimension 0 equal to number of bands/ channels
    ndxi_target: str
        the type of metric to calculate, by default "ndvi"
        if the metric is not in the predifined list, please specify "expr" and "params"
    working_dtype : DTypeLike
        data type used by the working intermediate numpy array, by default np.float32
    satellite_name: str
        currently only support sentinel-2, see 'ndxi_collection' for more information.
    expr: str, optional
        the expression used for metrics calculation
        required if the metric is not defined in ndxi_collection's sub-dictionary (e.g. S2_ndxi_parameters)
    params: List[str], optional
        the parameters used in the expression, e.g., ["a", "b"] is used in the expression "(a+b)/(a-b)"
        required if expr is not None
    nda_mask : np.ndarray, optional
        A numpy array for the valid part of the data, by default None.
        If provided, assumed to be in the same shape as nda[0, :, :]
    Returns
    -------
    Tuple[np.ndarray, np.ndarray]
        The first one is the calculated ndarray for the ndxi.
        The second one is the calculated mask array for the valid part.
    """
    nda = nda.astype(working_dtype)
    selected_expr, local_dict = None, None
    # find valid sr value range (for S2, valid range is 0-10000)
    if satellite_name in satellites_sr_valid_range_dict.keys():
        if satellites_sr_valid_range_dict[satellite_name] is None:
            raise ValueError(
                f"satellites_sr_valid_range_dict's {satellite_name} parameter. needs modifications."
            )
        else:
            satellite_sr_valid_range = satellites_sr_valid_range_dict[satellite_name]
    # check if satellite_name and ndxi_target are supported
    # for all other cases, selected_expr and local_dict are None
    if satellite_name in ndxi_collection.keys():
        if ndxi_collection[satellite_name] is None:
            raise ValueError(
                f"ndxi_collection's {satellite_name} parameter. needs modifications."
            )
        elif ndxi_target in ndxi_collection[satellite_name].keys():
            local_dict = {
                # SR input must be normalized to 0~1 for NDXI that have adjustment factor (e.g. EVI)
                param: nda[i] / satellite_sr_valid_range[2]
                for i, param in enumerate(
                    ndxi_collection[satellite_name][ndxi_target]["params"]
                )
            }
            selected_expr = ndxi_collection[satellite_name][ndxi_target]["expr"]
    # if selected_expr and local_dict are None, use user defined operation.
    if (selected_expr is None) and (local_dict is None):
        if not expr or not params:
            raise NotImplementedError(
                f"derived.py: yet to support {satellite_name} paramters."
            )
        else:
            local_dict = {
                # SR input must be normalized to 0~1 for NDXI that have adjustment factor (e.g. EVI)
                param: nda[i] / satellite_sr_valid_range[2]
                for i, param in enumerate(params)
            }
            selected_expr = expr
    # compute ndxi with ne
    with np.errstate(divide="ignore", invalid="ignore"):
        nda_ndxi = ne.evaluate(selected_expr, local_dict=local_dict)

    # compute and combine mask, but does not mask out nda_ndxi
    nda_tmp = np.isinf(nda_ndxi)
    nda_tmp |= np.isnan(nda_ndxi)
    # compute mask given input band
    for param in local_dict.keys():
        nda_tmp |= np.logical_or(
            local_dict[param]
            < (satellite_sr_valid_range[0] / satellite_sr_valid_range[2]),
            local_dict[param]
            > (satellite_sr_valid_range[1] / satellite_sr_valid_range[2]),
        )
    if nda_mask is None:
        nda_mask = ~nda_tmp
    else:
        nda_mask = nda_mask & (~nda_tmp)
    # NOTE: if True in this nda_mask, value should be kept.

    return nda_ndxi, nda_mask


def compute_mask(
    nda_QCs: np.ndarray,
    mask_index_lst: Optional[list] = [0],
    VALS_TO_MASK: list[Optional[list]] = [[0, 3, 8, 9, 11]],
    RANGE_NOT_TO_MASK: list[Optional[tuple[float, float]]] = [None],
    target_dimension: Optional[tuple[int, int]] = None,
    resample_mask: bool = False,
    verbose_lv: int = 0,
):
    """compute mask from 3D QCs numpy array

    Parameters
    ----------
    nda_QCs : np.ndarray
        a 3D ndarray where dimension 0 are bands, dimension 1,2 are x,y. This is used for computing mask.
    mask_index_lst : Optional[list], optional
        a list to indicate which bands of rds_QCs will be processed.
        For example, if [0,1] is provided, this will loop through band 0 and 1 and produce a combined mask.
        by default [0] to process band 0.
    VALS_TO_MASK : list[Optional[list]], optional
        a list of list to instruct which value to mask out.
        The length of the outer list should match the length of mask_index_lst.
        The inner list can be a list or None.
        For example, if [[0, 3, 8], None] is provided,
        pixels with value of 0, 3, or 8 in band 0 are masked out while doing nothing with band 1.
        by default [[0, 3, 8, 9, 11]]
    RANGE_NOT_TO_MASK : list[Optional[tuple[float, float]]], optional
        a list of list to instruct which value range not to mask out.
        The length of the outer list should match the length of mask_index_lst.
        The inner list can be a list/tuple of two numbers or None.
        For example, if [None, [50,100]] is provided,
        it will do nothing with band 0 while mask out pixels with values>100 or values<50 in band 1.
        by default [None]
    target_dimension: Optional[np.ndarray], optional
        Check if mask's x, y dimension matches user's expectation, e.g. (10980, 10980), by default None.
        If None, will not check dimension.
    resample_mask : bool, optional
        whether to resample mask to match target_dimension using cv2.INTER_NEAREST interpolation method
        by default False
    verbose_lv : int, optional
        0: do not report status
        1: report status for test-run purposes.

    Returns
    -------
    np.ndarray
        mask, a ndarray where True indicates raster value should be assigned as invalid value.

    """
    # check if input is 3D array
    if len(np.shape(nda_QCs)) != 3:
        raise ValueError("compute_mask: input nda_QCs must be a 3d numpy array.")
    if np.shape(nda_QCs)[0] != len(mask_index_lst):
        raise ValueError(
            "compute_mask:"
            + "check length of mask_index_lst. It should matches the number of bands (dimension 0) in nda_QCs"
        )
    # check if instruction are correctly formatted
    if (len(mask_index_lst) != len(VALS_TO_MASK)) or (
        len(mask_index_lst) != len(RANGE_NOT_TO_MASK)
    ):
        raise ValueError(
            "compute_mask: len of mask_index_lst does not match VALS_TO_MASK or RANGE_NOT_TO_MASK"
        )
    # compute mask
    output_mask = None
    for mask_var_i, mask_index in enumerate(mask_index_lst):
        # loop through different pre-process task for each layer
        for task in ["mask_val", "mask_out_of_valid_range"]:
            output_mask_temp = None
            if (task == "mask_val") & (VALS_TO_MASK[mask_var_i] is not None):
                if verbose_lv > 0:
                    logger.info(
                        f"generating mask for mask nda band {mask_index}. VALS_TO_MASK = {VALS_TO_MASK[mask_var_i]}"
                    )
                output_mask_temp = np.isin(
                    nda_QCs[mask_index, :, :], VALS_TO_MASK[mask_var_i]
                )
            if (task == "mask_out_of_valid_range") & (
                RANGE_NOT_TO_MASK[mask_var_i] is not None
            ):
                if verbose_lv > 0:
                    logger.info(
                        f"generating mask for mask nda band {mask_index}."
                        + f"RANGE_NOT_TO_MASK = {RANGE_NOT_TO_MASK[mask_var_i]}"
                    )
                output_mask_temp = np.logical_or(
                    nda_QCs[mask_index, :, :] < RANGE_NOT_TO_MASK[mask_var_i][0],
                    nda_QCs[mask_index, :, :] > RANGE_NOT_TO_MASK[mask_var_i][1],
                )
            if output_mask is None:
                output_mask = output_mask_temp
            else:
                output_mask = np.logical_or(output_mask, output_mask_temp)
    # resample mask if output mask dimension does not match expectation,
    if target_dimension is not None:
        if target_dimension != np.shape(output_mask)[-2:]:
            if resample_mask:
                output_mask = cv2.resize(
                    output_mask.astype(np.uint8),
                    target_dimension,
                    interpolation=cv2.INTER_NEAREST,
                )
            else:
                raise ValueError(
                    "compute_mask, "
                    + f"output_mask shape ({np.shape(output_mask)}) != target_dimension ({target_dimension})"
                )
    # NOTE: if True in this mask, value should be replaced
    return output_mask


def compute_masked_ndxi(
    rds_SRs: RasterDataset = None,
    rds_QCs: RasterDataset = None,
    ndxi_target_lst: List[str] = ["ndvi"],
    band_lst: Optional[List] = None,
    satellite_name: str = "Sentinel-2",
    mask_mode: str = "L3A",
    mask_filler_value: float = np.nan,
    expr: Optional[str] = None,
    params: List[str] = None,
    computation_mode: str = None,
):
    """compute masked ndxi, given raster datasets of Surface reflectance and QC layers.
    See skywalker-docs/tutorials/tutorial_utils_ndxi.ipynb for usage example.

    Parameters
    ----------
    rds_SRs : RasterDataset, optional
        raster dataset of n bands, used for computing ndxi, by default None
    rds_QCs : RasterDataset, optional
        raster dataset of n bands, used for computing mask, by default None
    ndxi_target_lst : list, optional
        which ndxi to process , by default ["ndvi"]
    band_lst: List, optional
        if compute multiple ndxis, need to provide the band list that is exactly the same order when retrieve
        raster items
    satellite_name: str
        currently only support sentinel-2, see 'ndxi_collection' for more information.
    mask_mode : str, optional
        define how mask should be processed, refer to mask_dict in common/rs_utils/derived.py, by default "L3A"
    mask_filler_value: float, optional
        customize value to fill-in masked data.
        by default np.nan
    expr: str, optional
        the expression used for metrics calculation
        required if satellite_name is not defined in ndxi_collection or
        if ndxi_target is not defined in ndxi_collection's sub-dictionary (e.g. S2_ndxi_parameters)
    params: List[str], optional
        the parameters used in the expression, e.g., ["a", "b"] is used in the expression "(a+b)/(a-b)"
        required if expr is not None.
    computation_mode : (deprecated) str, optional
        define how ndxi will be computed.
        will support more comprehensive method (e.g. allow 3 bands ndxi) in the future


    Returns
    -------
    Dict
       a dict contains 2d numpy arrays of ndxis, masked by QCs
    """
    # assume nda_QCs is stack according mask_dict["mask_mode"]["bands"], e.g. for L3A ["SCL", "QA"]
    nda_mask = compute_mask(
        nda_QCs=rds_QCs.data,
        mask_index_lst=mask_dict[satellite_name][mask_mode]["mask_index"],
        VALS_TO_MASK=mask_dict[satellite_name][mask_mode]["VALS_TO_MASK"],
        RANGE_NOT_TO_MASK=mask_dict[satellite_name][mask_mode]["RANGE_NOT_TO_MASK"],
        target_dimension=rds_SRs.data[0, :, :].shape,
        resample_mask=False,
    )
    # compute ndxi
    # assume rds_SRs is stack according to S2_ndxi_parameters["ndxi"]["bands"], e.g. for ndvi ["B04", "B08"]
    # where compute_ndxi returns invalid values.
    if computation_mode:
        logger.warning(
            "the function is upgraded and no longer require to specify computation mode"
        )
    ndxi_dict = {ndxi: None for ndxi in ndxi_target_lst}
    if band_lst is None:
        logger.info(
            f"no band name list is provided, will retrieve from the built-in dictionary for {ndxi_target_lst}"
        )
        if len(ndxi_target_lst) > 1:
            raise Exception("only one metric is allowed in this case")
        band_lst = ndxi_collection[satellite_name][ndxi_target_lst[0]]["bands"]
    for ndxi_target in ndxi_target_lst:
        band_target = ndxi_collection[satellite_name][ndxi_target]["bands"]
        band_target_idx = [band_lst.index(band) for band in band_target]
        nda_ndxi, nda_ndxi_mask = compute_ndxi(
            nda=rds_SRs.data[band_target_idx, :, :],
            ndxi_target=ndxi_target,
            satellite_name=satellite_name,
            expr=expr,
            params=params,
        )
        # mask out invalid values as mask_filler_value
        nda_ndxi = np.where(
            nda_mask, mask_filler_value, nda_ndxi
        )  # True in nda_mask == filter out
        nda_ndxi = np.where(
            ~nda_ndxi_mask, mask_filler_value, nda_ndxi
        )  # True in inverted nda_ndxi_mask == filter out
        ndxi_dict[ndxi_target] = nda_ndxi
    return ndxi_dict
